# 1. 导入依赖库
import numpy as np
import pandas as pd
import plotly
import plotly.graph_objs as go
from utils.LinearRegression.LinearRegressionUtils import LinearRegression

# 2. 读取数据集
data = pd.read_csv('../static/data/world-happiness-report-2017.csv')

# 2.1 定义测试数据（0.2）和训练数据（0.8）
train_data = data.sample(frac=0.8, random_state=200)
test_data = data.drop(train_data.index)

# 2.2 整理数据，确定特征和label
input_param_name_1 = "Economy..GDP.per.Capita."
input_param_name_2 = "Freedom"
output_param_name = "Happiness.Score"
x_train = train_data[[input_param_name_1, input_param_name_2]].values
y_train = train_data[[output_param_name]].values
x_test = test_data[[input_param_name_1, input_param_name_2]].values
y_test = test_data[[output_param_name]].values

# 2.3 初步展示初始数据（3D）
# 2.3.1 定义3D散点轨迹图——训练数据
plot_train_trace = go.Scatter3d(
    # flatten：将二维数组转化为一维数组
    # x_train[:, 0]：取出第0列的所有元素
    x=x_train[:, 0].flatten(),
    y=x_train[:, 1].flatten(),
    z=y_train.flatten(),
    name='train data',
    # 显示数据点
    mode='markers',
    marker={
        # 数据点的大小
        'size': 10,
        # 颜色
        'color': '#ebb563',
        # 透明度
        'opacity': 1,
        # 边框样式
        'line': {
            # 边框颜色
            'color': 'rgba(255,255,255)',
            # 边框宽度
            'width': 1
        }
    }
)
# 2.3.2 定义3D散点轨迹图——测试数据
plot_test_trace = go.Scatter3d(
    x=x_test[:, 0].flatten(),
    y=x_test[:, 1].flatten(),
    z=y_test.flatten(),
    name='test data',
    mode='markers',
    marker={
        'size': 10,
        'color': 'gray',
        'opacity': 1,
        'line': {
            'color': 'rgba(255,255,255)',
            'width': 1
        }
    }
)
# 2.3.3 定义散点图的布局
layout = go.Layout(
    title='World Happiness Report 2017',
    scene=dict(
        xaxis=dict(title=input_param_name_1),
        yaxis=dict(title=input_param_name_2),
        zaxis=dict(title=output_param_name),
        aspectmode="cube"  # 保持坐标轴比例一致
    ),
    margin={'l': 20, 'r': 20, 'b': 20, 't': 20}
)
# 2.3.4 绘制3D散点图
plot_data = [plot_train_trace , plot_test_trace]
fig = go.Figure(data=plot_data, layout=layout)
plotly.offline.iplot(fig)

# 3. 训练数据
linear_regression = LinearRegression(y_train, x_train)
cost_history = linear_regression.train(0.01,500)

# 4. 输出损失值
print("Cost before training: ", cost_history[0])
print("Cost after training: ", cost_history[-1])

# 5. 绘制预测曲面图
# 5.1 定义预测的所有节点坐标(x,y,z)
prediction_num=100
x_prediction_min = x_train[:,0].min()
x_prediction_max = x_train[:,0].max()
y_prediction_min = x_train[:,1].min()
y_prediction_max = x_train[:,1].max()
# 定义x与y的矩阵大小
x_predictions = np.zeros((prediction_num*prediction_num,1))
y_predictions = np.zeros((prediction_num*prediction_num,1))
# 填充x_predictions与y_predictions矩阵
x_y_index=0
for x_value in np.linspace(x_prediction_min, x_prediction_max, prediction_num):
    for y_value in np.linspace(y_prediction_min, y_prediction_max, prediction_num):
        x_predictions[x_y_index] = x_value
        y_predictions[x_y_index] = y_value
        x_y_index+=1
# 通过之前训练的theta,计算预测值z
z_predictions = linear_regression.get_predict(np.hstack((x_predictions,y_predictions)))
# 5.2 绘制预测散点图
plot_prediction_trace = go.Scatter3d(
    x=x_predictions.flatten(),
    y=y_predictions.flatten(),
    z=z_predictions.flatten(),
    name='prediction data',
    mode='markers',
    marker={
        'size': 1,
        'color': 'green',
        'opacity': 1,
        'line': {
            'color': 'rgba(255,255,255)',
            'width': 1
        }
    },
    surfaceaxis=2,  # 绘制预测曲面图

)

plot_data_prediction = [plot_train_trace , plot_test_trace, plot_prediction_trace]
fig_prediction = go.Figure(data=plot_data_prediction, layout=layout)
plotly.offline.iplot(fig_prediction)



